#ifndef DCU_KERNEL_H
#define DCU_KERNEL_H

#include "matrix.h"
#include <string>
#include <vector>
#include <map>
#include <mutex>
#include <hip/hip_runtime.h>
#include <hipblas/hipblas.h>

// Keep used kernels
__global__ void matrixMultiplyKernel(const float* A, const float* B, float* C, int M, int N, int K);
__global__ void matrixMultiplySharedKernel(const float* A, const float* B, float* C, int M, int N, int K);

// DCU内存池，用于重用设备内存
class DCUMemoryPool {
private:
    static std::map<size_t, std::vector<float*>> freeBuffers;
    static std::mutex mtx;
    static bool initialized;

public:
    // 初始化内存池
    static void init();
    
    // 释放内存池
    static void cleanup();
    
    // 分配内存
    static float* allocate(size_t size);
    
    // 释放内存
    static void release(float* ptr, size_t size);
};

// DCU上矩阵的表示
class DCUMatrix {
private:
    int rows;
    int cols;
    float* deviceData; // DCU上的数据指针
    bool ownsData;     // 标记是否需要释放设备内存
    bool fromPool;     // 标记内存是否来自内存池

public:
    // 构造函数
    DCUMatrix(int rows, int cols);
    
    // 使用预分配的设备内存构造矩阵
    DCUMatrix(int rows, int cols, float* externalData, bool takeOwnership = false);
    
    // 从主机矩阵直接构造
    DCUMatrix(const Matrix& hostMatrix);
    
    // 析构函数
    ~DCUMatrix();
    
    // 从主机矩阵中复制数据到设备
    void copyFromHost(const Matrix& hostMatrix);
    
    // 从设备复制数据到主机矩阵
    void copyToHost(Matrix& hostMatrix) const;
    
    // 异步从主机矩阵中复制数据到设备
    void copyFromHostAsync(const Matrix& hostMatrix, hipStream_t stream);
    
    // 异步从设备复制数据到主机矩阵
    void copyToHostAsync(const Matrix& hostMatrix, hipStream_t stream) const;
    
    // 获取行数
    int getRows() const { return rows; }
    
    // 获取列数
    int getCols() const { return cols; }
    
    // 获取设备数据指针
    float* getDeviceData() { return deviceData; }
    const float* getDeviceData() const { return deviceData; }
};

// Keep Buffer struct (used by multiplyPipelined)
struct Buffer {
    float* res_warp;       // 结果缓冲区 (Host, Pinned, float*)
    float* d_res_warp;     // 设备端结果缓冲区 (Device, float*)
    hipStream_t stream;     // 关联的流
    
    void alloc(int size);   // 分配内存 (size in bytes)
    void dealloc();         // 释放内存
};

// DCU上的矩阵乘法接口
namespace DCUMatrixMultiply {
    // Keep core functions
    bool initDCU();
    void releaseDCU();
    int getDeviceCount();
    void printDeviceInfo(int deviceId = 0);
    void multiply(const Matrix& A, const Matrix& B, Matrix& C); // Main entry point
    void multiplyHipBLAS(const Matrix& A, const Matrix& B, Matrix& C);
    void multiplyShared(const Matrix& A, const Matrix& B, Matrix& C);
    void multiplyLargeMatrix(const Matrix& A, const Matrix& B, Matrix& C);
    void multiplyPipelined(const Matrix& A, const Matrix& B, Matrix& C);

    // Keep error checking utilities
    void checkError(hipError_t error, const char* file, int line);
    void checkBlasError(hipblasStatus_t error, const char* file, int line);
}

// DCU矩阵乘法工厂函数
void dcuMatrixMultiply(const Matrix& A, const Matrix& B, Matrix& C);

// 错误检查宏
#define HIP_CHECK(err) DCUMatrixMultiply::checkError(err, __FILE__, __LINE__)
#define HIPBLAS_CHECK(err) DCUMatrixMultiply::checkBlasError(err, __FILE__, __LINE__)

#endif // DCU_KERNEL_H 